-module (osmos).

-export ([ start/0,
	   stop/0,
	   open/2,
	   close/1,
	   read/2,
	   write/3,
	   select_range/5,
	   select_continue/3 ]).

-behaviour (application).
-export ([ start/2,
	   stop/1 ]).

-ifdef (HAVE_EUNIT).
-include_lib ("eunit/include/eunit.hrl").
-endif.

-include ("osmos.hrl").
-include ("osmos_format.hrl").

%
% public
%

%% @spec start () -> ok
%% @doc Start the osmos application.
%% @end
start () ->
  application:start (osmos).

%% @spec stop () -> ok
%% @doc Stop the osmos application.
%% @end
stop () ->
  application:stop (osmos).

%% @spec open (Table::atom(), [Option]) -> { ok, Table::atom() }
%%                                       | { error, Reason }
%%       Option = { directory, string() }
%%              | { format, format() }
%%              | { max_memory_entries, integer() }
%%              | { max_selects, integer() }
%%              | { select_timeout_secs, number() }
%% @doc Open or reopen a osmos table. Note that more than one process may
%% open a given table; it is not actually closed until every process has
%% closed it. If another process has already opened the same table with
%% incompatible options, the call fails with {error, options_mismatch}.
%%
%% If a table has been opened previously in the given directory, it is
%% reopened. Otherwise, a new, empty table is created.
%% @end
open (Table, Options) when is_atom (Table), is_list (Options) ->
  osmos_server:open (Table, Options).

%% @spec close (Table::atom()) -> ok | { error, Reason }
%% @doc Close the given table, which must have been opened by the calling
%% process. The table is not actually closed until all processes that
%% opened it have closed it.
%% @end
close (Table) when is_atom (Table) ->
  osmos_server:close (Table).

%% @spec read (atom(), any()) -> { ok, Value::any() } | not_found
%% @doc Read the value for Key in the given Table. Return { ok, Value } if
%% found, or not_found if not.
%% @end
read (Table, Key) when is_atom (Table) ->
  { ok, Pid } = osmos_server:get_pid (Table),
  osmos_table:read (Pid, Key).

%% @spec write (atom(), any(), any()) -> ok | { error, Reason }
%% @doc Write a new Value for Key in the given Table. The exact semantics
%% of writing to the table depend on the table's format; in particular,
%% writing may add a new value, update an existing value, or delete a value,
%% depending on the format's merge, short-circuit, and delete methods.
%% @end
write (Table, Key, Value) when is_atom (Table) ->
  { ok, Pid } = osmos_server:get_pid (Table),
  osmos_table:write (Pid, Key, Value).

%% @spec select_range (atom(), Less, Less, Select, integer())
%%         -> { ok, [ { Key::any(), Value::any() } ], continuation() }
%%          | { error, Reason }
%%       Less = (Key::any()) -> bool()
%%       Select = (Key::any(), Value::any()) -> bool()
%% @doc Begin iterating over a range of keys and values in the table
%% specified by lower and upper bounds, and a predicate that can select
%% within this range.
%%
%% The lower bound LessLo (Key) should return true iff Key is less than
%% every element in the range. The upper bound LessHi (Key) should return
%% true iff Key is not greater than every element in the range.
%% The selection predicate Select (Key, Value) should return true to
%% include {Key, Value} in the results.
%%
%% At most N key-value pairs are returned if successful, along with a
%% continuation that can be used to retrieve additional results with
%% select_continue. If fewer than N results are returned from select_range
%% or select_continue, the end of the table has been reached; further calls
%% to select_continue will return an empty list of results.
%%
%% The results from subsequent calls to select_continue always reflect
%% the state of the table at the time of the original call to select_range.
%% That is, any intervening writes to the table will not affect the
%% results from the select.
%%
%% Note that, depending on the select_timeout_secs and max_selects options
%% given when opening the table, the continuation may go stale (if more
%% results are not requested within select_timeout_secs, or if max_selects
%% other selects are started in the meantime).
%% @end
select_range (Table, LessLo, LessHi, Select, N) when is_atom (Table),
						     is_function (LessLo),
						     is_function (LessHi),
						     is_function (Select),
						     is_integer (N) ->
  { ok, Pid } = osmos_server:get_pid (Table),
  osmos_table:select_range (Pid, LessLo, LessHi, Select, N).

%% @spec select_continue (atom(), continuation(), integer())
%%         -> { ok, [ { Key::any(), Value::any() } ], continuation() }
%%          | { error, Reason }
%% @doc Continue selecting results using a continuation returned by
%% select_range or select_continue.
%% @end
select_continue (Table, Cont, N) when is_atom (Table), is_integer (N) ->
  { ok, Pid } = osmos_server:get_pid (Table),
  osmos_table:select_continue (Pid, Cont, N).

%
% application callbacks
%

start (_Type, _Args) ->
  osmos_supervisor:start_link ().

stop (_State) ->
  ok.

%
% tests
%

-ifdef (EUNIT).

make_setup () ->
  fun () ->
    { ok, [ App ] } = file:consult ("../src/osmos.app"),
    ok = application:load (App),
    ok = osmos:start ()
  end.

make_teardown (Dir) ->
  fun (_) ->
    osmos:stop (),
    application:unload (osmos),
    os:cmd ("rm -rf '" ++ Dir ++ "'")
  end.

test_tmp_path (Line) ->
  lists:flatten (io_lib:format ("tmp-test-~s-~s-~b",
				[ ?MODULE, os:getpid (), Line ])).

basic_test_ () ->
  Dir = test_tmp_path (?LINE),
  { setup,
    make_setup (),
    make_teardown (Dir),
    fun () ->
      Format = osmos_table_format:binary_replace (64),
      Options = [ { directory, Dir },
		  { format, Format },
		  { max_memory_entries, 4 } ],
      Table = foo,

      { ok, Table } = osmos:open (Table, Options),
      List1 = [ <<"foo1">>, <<"bar1">>, <<"baz1">>, <<"quux1">>, <<"blorf1">> ],
      ok = assert_write_md5 (?LINE, Table, List1),
      ok = assert_read_md5 (?LINE, Table, List1),
      ok = osmos:close (Table),

      { ok, Table } = osmos:open (Table, Options),
      ok = assert_read_md5 (?LINE, Table, List1),
      List2 = [ <<"foo2">>, <<"bar2">>, <<"baz2">>, <<"quux2">>, <<"blorf2">> ],
      ok = assert_write_md5 (?LINE, Table, List2),
      ok = assert_read_md5 (?LINE, Table, List2),
      ok = assert_read_md5 (?LINE, Table, List1),
      ok = osmos:close (Table),

      { ok, Table } = osmos:open (Table, Options),
      ok = assert_read_md5 (?LINE, Table, List1),
      ok = assert_read_md5 (?LINE, Table, List2),
      List3 = [ <<"foo3">>, <<"bar3">>, <<"baz3">>, <<"quux3">>, <<"blorf3">> ],
      ok = assert_write_md5 (?LINE, Table, List3),
      ok = assert_read_md5 (?LINE, Table, List1),
      ok = assert_read_md5 (?LINE, Table, List2),
      ok = assert_read_md5 (?LINE, Table, List3),
      timer:sleep (1000),	% wait for merge to complete
      ok = assert_read_md5 (?LINE, Table, List1),
      ok = assert_read_md5 (?LINE, Table, List2),
      ok = assert_read_md5 (?LINE, Table, List3),
      ok = osmos:close (Table),

      ok
    end
  }.

assert_write_md5 (Line, Table, [ Key | Keys ]) ->
  case osmos:write (Table, Key, erlang:md5 (Key)) of
    ok -> ok;
    Bad -> throw ({ write_error, Line, Key, Bad })
  end,
  assert_write_md5 (Line, Table, Keys);
assert_write_md5 (_Line, _Table, []) ->
  ok.

assert_read_md5 (Line, Table, [ Key | Keys ]) ->
  Expected = erlang:md5 (Key),
  case osmos:read (Table, Key) of
    { ok, Expected } -> ok;
    Bad -> throw ({ read_error, Line, Key, Bad })
  end,
  assert_read_md5 (Line, Table, Keys);
assert_read_md5 (_Line, _Table, []) ->
  ok.

medium_test_ () ->
  Dir = test_tmp_path (?LINE),
  { setup,
    make_setup (),
    make_teardown (Dir),
    { timeout,
      60,
      fun () ->
	Format = osmos_table_format:binary_replace (4096),
	Options = [ { directory, Dir },
		    { format, Format },
		    { max_memory_entries, 1024 } ],
	Table = foo,

	{ ok, Table } = osmos:open (Table, Options),
	ok = medium_test_read_write (Table, 1, 0, 10000),
	ok = osmos:close (Table),

	{ ok, Table } = osmos:open (Table, Options),
	ok = medium_test_read_write (Table, 1/2, 10000, 20000),
	ok = medium_test_read_write (Table, 1/3, 20000, 30000),
	ok = medium_test_read_write (Table, 1/4, 30000, 40000),
	ok = medium_test_read_write (Table, 0, 40000, 50000),
	ok = osmos:close (Table),

	ok
      end
    }
  }.

medium_test_read_write (Table, PRead, Start, Max) ->
  T0 = now (),
  ok = medium_test_loop (Table, PRead, Start, Max),
  D = timer:now_diff (now (), T0),
  error_logger:info_msg ("medium_test ~p r-w/sec with P(read)=~p (~b)~n",
			[ 1000000 * (Max - Start) / D, PRead, (Max - Start) ]),
  ok.

medium_test_loop (_Table, _PRead, Max, Max) ->
  ok;
medium_test_loop (Table, PRead, N, Max) ->
  WriteKey = medium_test_key (N),
  WriteValue = medium_test_value (WriteKey),
  case osmos:write (Table, WriteKey, WriteValue) of
    ok -> ok;
    BadWrite -> throw ({ write_error, ?LINE, WriteKey, BadWrite })
  end,
  case random:uniform () < PRead of
    false ->
      ok;
    true ->
      ReadN = random:uniform (N+1) - 1,
      ReadKey = medium_test_key (ReadN),
      ReadValue = medium_test_value (ReadKey),
      case osmos:read (Table, ReadKey) of
	{ ok, ReadValue } -> ok;
	BadRead -> throw ({ read_error, ?LINE, ReadKey, ReadValue, BadRead })
      end
  end,
  medium_test_loop (Table, PRead, N+1, Max).

medium_test_key (N) ->
  << N:32/big-unsigned-integer >>.

medium_test_value (Key) ->
  N = size (Key) - 1,
  << _:N/binary, _:4, Len:4 >> = Key,
  Suffix = list_to_binary (lists:duplicate (Len, 0)),
  << Key/binary, Suffix/binary >>.

select_test_ () ->
  Dir = test_tmp_path (?LINE),
  { setup,
    make_setup (),
    make_teardown (Dir),
    { timeout,
      60,
      fun () ->
	Format = osmos_table_format:binary_replace (4096),
	Options = [ { directory, Dir },
		    { format, Format },
		    { max_memory_entries, 1024 } ],
	Table = foo,
	{ ok, Table } = osmos:open (Table, Options),
	ok = medium_test_loop (Table, 0.0, 1, 100000),
	ok = select_test_select (Table, 1, 100000,     0,    137, 100),
	ok = select_test_select (Table, 1, 100000, 13427,  15118, 100),
	ok = select_test_select (Table, 1, 100000, 99997, 100002, 100),
	ok = osmos:close (Table),
	ok
      end
    }
  }.

select_test_select (Table, Min, Max, Lo, Hi, N) ->
  LoKey = medium_test_key (Lo),
  LessLo = fun (K) -> K < LoKey end,
  HiKey = medium_test_key (Hi),
  LessHi = fun (K) -> K < HiKey end,
  Select = fun (_K, _V) -> true end,
  SelectLo = case Lo < Min of
	       true -> Min;
	       false -> Lo
	     end,
  SelectHi = case Hi > Max of
	       true -> Max;
	       false -> Hi
	     end,
  LoPlusN = SelectLo + N,
  ChunkHi = case LoPlusN > SelectHi of
	      true -> SelectHi;
	      false -> LoPlusN
	    end,
  Expected = case ChunkHi > SelectLo of
	       true ->
		 lists:map (fun (X) ->
			      K = medium_test_key (X),
			      { K, medium_test_value (K) }
			    end,
			    lists:seq (SelectLo, ChunkHi - 1));
	       false ->
		 []
	     end,
  N0 = N div 2,
  N1 = N - N0,
  SplitLen = case N0 > length (Expected) of
	       true  -> length (Expected);
	       false -> N0
	     end,
  { Expected0, Expected1 } = lists:split (SplitLen, Expected),
  case osmos:select_range (Table, LessLo, LessHi, Select, N0) of
    { ok, Entries0, Cont } ->
      case Entries0 =:= Expected0 of
	false ->
	  throw ({ select_wrong, ?LINE, Expected0, Entries0 });
	true ->
	  case osmos:select_continue (Table, Cont, N1) of
	    { ok, Entries1, _C } ->
	      case Entries1 =:= Expected1 of
		false ->
		  throw ({ select_wrong, ?LINE, Expected1, Entries1 });
		true ->
		  ok
	      end;
	    Bad1 ->
	      throw ({ select_error, ?LINE, Bad1 })
	  end
      end;
    Bad0 ->
      throw ({ select_error, ?LINE, Bad0 })
  end.

sum_test_ () ->
  Dir = test_tmp_path (?LINE),
  { setup,
    make_setup (),
    make_teardown (Dir),
    { timeout,
      300,
      fun () ->
	Format = osmos_table_format:term_sum_uint63_delete (4096),
	Options = [ { directory, Dir },
		    { format, Format },
		    { max_memory_entries, 8192 } ],
	Table = foo,
	{ ok, Table } = osmos:open (Table, Options),
	ok = sum_test_write_loop (Table, 10000),
	ok = sum_test_read_loop (Table, 10000),
	ok = osmos:close (Table),
	ok
      end
    }
  }.

sum_test_write_loop (Table, MaxN) ->
  sum_test_write_loop (Table, 1, 0, MaxN).
sum_test_write_loop (_Table, Pass, _N, _MaxN) when Pass > 64 ->
  ok;
sum_test_write_loop (Table, Pass, MaxN, MaxN) ->
  sum_test_write_loop (Table, Pass + 1, 0, MaxN);
sum_test_write_loop (Table, Pass, N, MaxN) ->
  WriteKey = sum_test_key (N),
  WriteValue = sum_test_write_value (N, Pass),
  case osmos:write (Table, WriteKey, WriteValue) of
    ok -> ok;
    BadWrite -> throw ({ write_error, ?LINE, WriteKey, BadWrite })
  end,
  sum_test_write_loop (Table, Pass, N + 1, MaxN).

sum_test_read_loop (Table, MaxN) ->
  sum_test_read_loop (Table, 0, MaxN).
sum_test_read_loop (_Table, MaxN, MaxN) ->
  ok;
sum_test_read_loop (Table, N, MaxN) ->
  ReadKey = sum_test_key (N),
  Got = osmos:read (Table, ReadKey),
  case sum_test_read_value (N) of
    delete ->
      case Got of
	not_found -> ok;
	_ -> throw ({ read_error, ?LINE, ReadKey, Got, delete })
      end;
    Expected ->
      case Got of
	{ ok, Expected } -> ok;
	_ -> throw ({ read_error, ?LINE, ReadKey, Got, Expected })
      end
  end,
  sum_test_read_loop (Table, N + 1, MaxN).

sum_test_key (N) ->
  Md5 = << X:4, _:4, _/binary >> = erlang:md5 (term_to_binary (N)),
  << Y:X/binary, _/binary >> = Md5,
  [ N | binary_to_list (Y) ].

sum_test_read_value (N) when is_integer (N) ->
  List = binary_split_2 (erlang:md5 (term_to_binary (N))),
  sum_test_read_value (N, delete, List).
sum_test_read_value (_N, X, []) ->
  X;
sum_test_read_value (N, _, [ 0 | Rest ]) ->
  sum_test_read_value (N, delete, Rest);
sum_test_read_value (N, delete, [ B | Rest ]) ->
  sum_test_read_value (N, B * N, Rest);
sum_test_read_value (N, A, [ B | Rest ]) ->
  sum_test_read_value (N, A + B * N, Rest).

sum_test_write_value (N, Pass) when Pass >= 1, Pass =< 64 ->
  List = binary_split_2 (erlang:md5 (term_to_binary (N))),
  case lists:nth (Pass, List) of
    0 -> delete;
    B when is_integer (B) -> B * N
  end.

binary_split_2 (B) ->
  binary_split_2 (B, []).
binary_split_2 (<<>>, Acc) ->
  lists:reverse (Acc);
binary_split_2 (<<A:2, B:2, C:2, D:2, Rest/binary>>, Acc) ->
  binary_split_2 (Rest, [ D, C, B, A | Acc ]).

vector_test_ () ->
  Dir = test_tmp_path (?LINE),
  { setup,
    make_setup (),
    make_teardown (Dir),
    { timeout,
      600,
      fun () ->
	Format = osmos_table_format:binary_sum_uint64_vector (4096),
	Options = [ { directory, Dir },
		    { format, Format },
		    { max_memory_entries, 1024 } ],
	Table = foo,
	{ ok, Table } = osmos:open (Table, Options),
	T0 = now (),
	NWrites = 1000000,
	Tree = vector_test_write_loop (Table, NWrites),
	DT = 1.0e-6 * timer:now_diff (now (), T0),
	error_logger:info_msg ("~b writes in ~p seconds, ~p writes/second~n",
			       [ NWrites, DT, NWrites / DT ]),
	ok = vector_test_verify (Table, Tree),
	ok = osmos:close (Table),
	ok
      end
    }
  }.

vector_test_write_loop (Table, N) ->
  vector_test_write_loop (Table, N, gb_trees:empty ()).

vector_test_write_loop (Table, N, Tree) when N > 0 ->
  KeyN = round (10000 * poisson ()),
  Remember = (KeyN < 10000 andalso (KeyN rem 7 =:= 0)),
  Key = erlang:md5 (term_to_binary (KeyN)),
  Vector = list_to_tuple ([ random:uniform (2) - 1 || _ <- lists:seq (1, 3) ]),
  NewTree = case Remember of
	      false ->
		Tree;
	      true ->
		case gb_trees:lookup (Key, Tree) of
		  { value, OldVector } ->
		    NewVector = sum_tuples (Vector, OldVector),
		    gb_trees:update (Key, NewVector, Tree);
		  none ->
		    gb_trees:insert (Key, Vector, Tree)
		end
	    end,
  ok = osmos:write (Table, Key, Vector),
  vector_test_write_loop (Table, N - 1, NewTree);
vector_test_write_loop (_Table, 0, Tree) ->
  Tree.

vector_test_verify (Table, Tree) ->
  vector_test_verify_loop (Table, gb_trees:iterator (Tree)).

vector_test_verify_loop (Table, It) ->
  case gb_trees:next (It) of
    { K, V, NewIt } ->
      case osmos:read (Table, K) of
	{ ok, V } ->
	  vector_test_verify_loop (Table, NewIt);
	{ ok, Other } ->
	  erlang:error ({ mismatch, K, V, Other });
	not_found ->
	  erlang:error ({ not_found, K, V })
      end;
    none ->
      ok
  end.

sum_tuples (T1, T2) ->
  Z = lists:zip (tuple_to_list (T1), tuple_to_list (T2)),
  list_to_tuple ([ N1 + N2 || { N1, N2 } <- Z ]).

poisson () ->
  X = random:uniform (),
  case X == 0 of
    true ->
      poisson ();
    false ->
      -math:log (X)
  end.

-endif.
